# -*- coding: UTF-8 -*-
"""
@项目名称：get_pocket.py
@作   者：陆地起飞全靠浪
@创建日期：2025-09-12-16:02
"""
import os
import numpy as np
from prody import parsePDB, writePDB
from rdkit import Chem


def get_mol2_xyz_from_cmd(ligand_mol2):
    x = os.popen("cat %s | sed -n '/@<TRIPOS>ATOM/,/@<TRIPOS>BOND/'p | awk '{print $3}'" % ligand_mol2).read().splitlines()[1:-1]
    y = os.popen("cat %s | sed -n '/@<TRIPOS>ATOM/,/@<TRIPOS>BOND/'p | awk '{print $4}'" % ligand_mol2).read().splitlines()[1:-1]
    z = os.popen("cat %s | sed -n '/@<TRIPOS>ATOM/,/@<TRIPOS>BOND/'p | awk '{print $5}'" % ligand_mol2).read().splitlines()[1:-1]
    return np.asanyarray(list(zip(x, y, z))).astype(float)


def get_pocket(reference_ligand_sdf, protein_pdb, reference_ligand_mol2, pocket_pdb, exwithin=12):
    # select & save pocket
    ligand_mol = Chem.MolFromMolFile(reference_ligand_sdf, removeHs=False)
    if ligand_mol is None:
        ligand_mol = Chem.MolFromMol2File(reference_ligand_mol2, removeHs=False)
    if ligand_mol is None:
        try:
            ligpos = get_mol2_xyz_from_cmd(reference_ligand_mol2)
        except:
            print(protein_pdb)
    else:
        ligpos = ligand_mol.GetConformer().GetPositions()
    if not os.path.exists(pocket_pdb) or os.path.getsize(pocket_pdb) == 0:
        protein_prody_obj = parsePDB(protein_pdb)
        condition = f'same residue as exwithin {exwithin} of somepoint'
        pocket_selected = protein_prody_obj.select(condition, somepoint=ligpos)  # （n, 3）
        writePDB(pocket_pdb, atoms=pocket_selected)


if __name__ == '__main__':
    # src
    reference_ligand_sdf = 'lig_4yzm_protein_src_1.sdf'
    protein_pdb = '4yzm_protein.pdb'
    reference_ligand_mol2 = 'lig_4yzm_protein_src_1.mol2'
    # dst
    pocket_pdb = 'pocket_ligH12A.pdb'
    get_pocket(reference_ligand_sdf, protein_pdb, reference_ligand_mol2, pocket_pdb)
